package com.mainbo.core.util;

import com.mainbo.core.common.exception.RsaSignException;

import java.security.InvalidKeyException;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.security.Signature;
import java.security.SignatureException;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;

/**
 * rsa 加密工具类
 * @author moshang
 * @date 2020-03-01
 **/
public class RSAUtils {
    public static final String SIGNATURE_ALGORITHM = "MD5withRSA";
    public static final SecureRandom DEFAULT_SECURE_RANDOM;
    private static final String SIGNATUR_NAME = "RSA";

    static {
        DEFAULT_SECURE_RANDOM = new SecureRandom();
        DEFAULT_SECURE_RANDOM.nextBytes(new byte[64]);
    }

    /**
     *
     * @param keySizeBits
     *          key size
     * @return key pair
     */
    public static KeyPair createRSAKeyPair(int keySizeBits) {
        KeyPairGenerator keyGenerator;
        try {
            keyGenerator = KeyPairGenerator.getInstance(SIGNATUR_NAME);
        } catch (NoSuchAlgorithmException e) {
            throw new IllegalStateException(
                    "Unable to obtain an RSA KeyPairGenerator: " + e.getMessage(), e);
        }

        keyGenerator.initialize(keySizeBits, DEFAULT_SECURE_RANDOM);

        return keyGenerator.genKeyPair();
    }

    /**
     *
     * @param keyPair
     *          keypair
     * @return private key
     */
    public static String getRSAPrivateKey(KeyPair keyPair) {
        return Encodes.encodeBase64(keyPair.getPrivate().getEncoded());
    }

    /**
     *
     * @param keyPair
     *          key pair
     * @return private key
     */
    public static String getRSAPublicKey(KeyPair keyPair) {
        return Encodes.encodeBase64(keyPair.getPublic().getEncoded());
    }

    /**
     * 得到公钥
     *
     * @param key
     *          密钥字符串（经过base64编码）
     * @throws Exception
     *           e
     */
    public static PublicKey getPublicKey(String key) {
        byte[] keyBytes = Encodes.decodeBase64(key);
        X509EncodedKeySpec keySpec = new X509EncodedKeySpec(keyBytes);
        PublicKey publicKey;
        try {
            KeyFactory keyFactory = KeyFactory.getInstance(SIGNATUR_NAME);
            publicKey = keyFactory.generatePublic(keySpec);
        } catch (NoSuchAlgorithmException | InvalidKeySpecException e) {
            throw new RsaSignException(e);
        }

        return publicKey;
    }

    /**
     * 得到私钥
     *
     * @param key
     *          密钥字符串（经过base64编码）
     * @throws Exception
     *           e
     */
    public static PrivateKey getPrivateKey(String key) {
        byte[] keyBytes;
        keyBytes = Encodes.decodeBase64(key);
        PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(keyBytes);
        PrivateKey privateKey;
        try {
            KeyFactory keyFactory = KeyFactory.getInstance(SIGNATUR_NAME);
            privateKey = keyFactory.generatePrivate(keySpec);
        } catch (NoSuchAlgorithmException | InvalidKeySpecException e) {
            throw new RsaSignException(e);
        }
        return privateKey;
    }

    /**
     *
     * @param privateKey
     *          密钥字符串（经过base64编码）
     * @param data
     *          要签名数据
     * @return 加密后数据
     */
    public static byte[] sign(String privateKey, byte[] data) {
        PrivateKey priK = getPrivateKey(privateKey);
        try {
            Signature sig = Signature.getInstance(SIGNATURE_ALGORITHM);
            sig.initSign(priK);
            sig.update(data);
            return sig.sign();
        } catch (SignatureException | InvalidKeyException
                | NoSuchAlgorithmException e) {
            throw new RsaSignException(e);
        }
    }

    /**
     *
     * @param pulicKey
     *          公匙
     * @param data
     *          要验证的数据
     * @param sign
     *          签名
     * @return true if success
     */
    public static boolean verify(String pulicKey, byte[] data, byte[] sign) {
        PublicKey pubK = getPublicKey(pulicKey);
        try {
            Signature sig = Signature.getInstance(SIGNATURE_ALGORITHM);
            sig.initVerify(pubK);
            sig.update(data);
            return sig.verify(sign);
        } catch (InvalidKeyException | NoSuchAlgorithmException
                | SignatureException e) {
            throw new RsaSignException(e);
        }
    }

    /**
     *
     * @param pulicKey
     *          公匙 String
     * @param plaintext
     *          to encrypt data
     * @return encrypt string
     */
    public static byte[] encrypt(String pulicKey, byte[] plaintext) {
        PublicKey publicKey = getPublicKey(pulicKey);
        return encrypt(publicKey, plaintext);
    }

    /**
     * 使用公匙加密
     *
     * @param publicKey
     *          公匙
     * @param plaintext
     *          to encrypt data
     * @return encrypt string
     */
    public static byte[] encrypt(PublicKey publicKey, byte[] plaintext) {
        try {
            Cipher cipher = Cipher.getInstance(SIGNATUR_NAME);
            cipher.init(Cipher.ENCRYPT_MODE, publicKey);
            return cipher.doFinal(plaintext);
        } catch (InvalidKeyException | NoSuchAlgorithmException
                | NoSuchPaddingException | IllegalBlockSizeException
                | BadPaddingException e) {
            throw new RsaSignException(e);
        }
    }

    /**
     * 使用私匙解密
     *
     * @param privateKey 私匙
     * @param encrypted 密文
     * @return 解密后数据
     */
    public static byte[] decrypt(String privateKey, byte[] encrypted) {
        PrivateKey pKey = getPrivateKey(privateKey);
        return decrypt(pKey, encrypted);
    }

    /**
     * 使用私匙解密
     *
     * @param pKey 私匙
     * @param encrypted 密文
     * @return 解密后数据
     */
    public static byte[] decrypt(PrivateKey pKey, byte[] encrypted) {
        try {
            Cipher cipher = Cipher.getInstance(SIGNATUR_NAME);
            cipher.init(Cipher.DECRYPT_MODE, pKey);
            return cipher.doFinal(encrypted);
        } catch (InvalidKeyException | NoSuchAlgorithmException
                | NoSuchPaddingException | IllegalBlockSizeException
                | BadPaddingException e) {
            throw new RsaSignException(e);
        }
    }

    /**
     * 使用私匙加密
     * @param privateKey 私匙
     * @param plaintext 数据
     * @return 加密后数据
     */
    public static byte[] encryptByPrivateKey(String privateKey,
                                             byte[] plaintext) {
        PublicKey publicKey = getPublicKey(privateKey);
        return encrypt(publicKey, plaintext);
    }

    /**
     * 使用私匙加密
     * @param privateKey 私匙
     * @param plaintext 数据
     * @return 加密后数据
     */
    public static byte[] encryptByPrivateKey(PrivateKey privateKey,
                                             byte[] plaintext) {
        try {
            Cipher cipher = Cipher.getInstance(SIGNATUR_NAME);
            cipher.init(Cipher.ENCRYPT_MODE, privateKey);
            return cipher.doFinal(plaintext);
        } catch (InvalidKeyException | NoSuchAlgorithmException
                | NoSuchPaddingException | IllegalBlockSizeException
                | BadPaddingException e) {
            throw new RsaSignException(e);
        }
    }

    /**
     * 使用公匙解密
     * @param pubKey 公匙
     * @param encrypted 密文
     * @return 解密后数据
     */
    public static byte[] decryptByPublicKey(String pubKey, byte[] encrypted) {
        PrivateKey pKey = getPrivateKey(pubKey);
        return decrypt(pKey, encrypted);
    }

    /**
     * 使用公匙解密
     * @param pubKey 公匙
     * @param encrypted 密文
     * @return 解密后数据
     */
    public static byte[] decryptByPublicKey(PublicKey pubKey, byte[] encrypted) {
        try {
            Cipher cipher = Cipher.getInstance(SIGNATUR_NAME);
            cipher.init(Cipher.DECRYPT_MODE, pubKey);
            return cipher.doFinal(encrypted);
        } catch (InvalidKeyException | NoSuchAlgorithmException
                | NoSuchPaddingException | IllegalBlockSizeException
                | BadPaddingException e) {
            throw new RsaSignException(e);
        }
    }
}
